from __future__ import annotations

import json
from dataclasses import dataclass, field, fields
from typing import List, Optional, Annotated, Any, Union, Dict

from langchain_core.messages import AnyMessage, BaseMessage
from langgraph.graph import add_messages
from langchain_core.runnables import RunnableConfig, ensure_config, AddableDict

from ..config.agents import AGENT_FINAL_ANSWER_STR, get_action_final_answer


@dataclass(kw_only=True)
class GraphNodeParam:
    type: str = None
    params: dict = field(default_factory=dict)
    kwargs: dict = field(default_factory=dict)

    def add_param(self, key: str, value: Any):
        self.params[key] = value

    def get_param(self, key: str, default: Any = None):
        return self.params[key] if self.params[key] else default

    def __str__(self) -> str:
        return json.dumps(self.to_dict(), ensure_ascii=False)

    def to_dict(self) -> dict:
        return {
            "type": self.type,
            "params": self.params,
            "kwargs": self.kwargs
        }

    def __json__(self) -> dict:
        return self.to_dict()


@dataclass(kw_only=True)
class Usage:
    model: str
    input_tokens: int = 0
    output_tokens: int = 0
    total_tokens: int = 0

    def add_tokens(self, input_tokens: int, output_tokens: int):
        self.input_tokens += input_tokens
        self.output_tokens += output_tokens
        self.total_tokens = self.input_tokens + self.output_tokens

    def to_dict(self) -> dict:
        return {
            "model": self.model,
            "input_tokens": self.input_tokens,
            "output_tokens": self.output_tokens,
            "total_tokens": self.total_tokens
        }

    def __json__(self) -> dict:
        return self.to_dict()


@dataclass(kw_only=True)
class GraphState:
    messages: Annotated[List[BaseMessage], add_messages] = field(default_factory=list)
    loop_nodes: List[GraphNodeParam] = field(default_factory=list)
    goto: GraphNodeParam = field(default_factory=GraphNodeParam)
    force_goto: GraphNodeParam = field(default_factory=GraphNodeParam)
    loop_retry: int = 0
    usages: List[Usage] = field(default_factory=list)
    prefix_executed: bool = False
    is_task_done: bool = False
    ext: dict = field(default_factory=dict)

    def add_usage(self, model: str, input_tokens: int, output_tokens: int):
        for usage in self.usages:
            if usage.model == model:
                usage.add_tokens(input_tokens, output_tokens)
                return

        new_usage = Usage(
            model=model,
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            total_tokens=input_tokens + output_tokens
        )
        self.usages.append(new_usage)

    def get_usages(self):
        if not self.usages:
            return []
        return list(map(lambda x: x.__dict__, self.usages))


@dataclass(kw_only=True)
class GraphRequest:
    question: str = ""
    image: str | List = ""
    thread_id: str = field()
    history: List[AnyMessage] = field(default_factory=list)
    kwargs: dict = field(default_factory=dict)
    context: str = ""
    mcp_tools: str = ""
    web_search: bool = True
    agent_name: str = None

    @classmethod
    def from_runnable_config(
        cls, config: Optional[RunnableConfig] = None
    ) -> GraphRequest:
        """Load configuration w/ defaults for the given invocation."""
        config = ensure_config(config)
        configurable = config.get("configurable") or {}
        _fields = {f.name for f in fields(cls) if f.init}
        return cls(**{k: v for k, v in configurable.items() if k in _fields})


@dataclass
class AgentStepLog:
    output: Union[str, List[Dict[str, Any]], GraphNodeParam] = None
    meta: Dict[str, Any] = field(default_factory=dict)
    type: str = "answer"

    @classmethod
    def build_answer(cls, chunk: Union[AddableDict, Dict, str], meta: Dict[str, Any] = {}) -> "AgentStepLog":
        chunk = "" if not chunk else chunk
        if (isinstance(chunk, AddableDict) or isinstance(chunk, Dict)) and "output" in chunk:
            output = chunk["output"]
            if AGENT_FINAL_ANSWER_STR in output:
                output = get_action_final_answer(output)
            return cls(output=output, meta=meta)
        elif isinstance(chunk, str):
            return cls(output=chunk, meta=meta)

    @classmethod
    def build_tool_action(cls, action: str, output: Any = "", meta: Dict[str, Any] = None) -> "AgentStepLog":
        meta = meta or {}
        meta["action"] = action
        return cls(output=output, type="tool", meta=meta)

    @classmethod
    def build_log_action(cls, action: str, output: Any = "", meta: Dict[str, Any] = None) -> "AgentStepLog":
        meta = meta or {}
        meta["action"] = action
        return cls(output=output, type="log", meta=meta)

    def __str__(self):
        output_dict = self.__dict__.copy()
        if isinstance(output_dict["output"], GraphNodeParam):
            output_dict["output"] = output_dict["output"].to_dict()
        return json.dumps(output_dict, indent=2, ensure_ascii=False)
